import numpy as np
from typing import List, Tuple
import logging
from datetime import datetime
import sys
import os

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..'))
if project_root not in sys.path:
    sys.path.insert(0, project_root)
# =============================================

try:
    from config import settings
    from utils.logger import setup_logger
except ImportError as e:
    print(f"导入错误: {e}")
    print("正在使用备用方案...")


    # 备用方案：直接在这里定义必要的类和函数
    class Settings:
        NAME_SIMILARITY_THRESHOLD = 0.8
        SPEC_SIMILARITY_THRESHOLD = 0.5
        NAME_SIM_WEIGHT = 0.5
        SPEC_SIM_WEIGHT = 0.5


    settings = Settings()


    def setup_logger(name):
        import logging
        logger = logging.getLogger(name)
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            logger.setLevel(logging.INFO)
        return logger

logger = setup_logger(__name__)


class NumpySearcher:
    """基于NumPy的向量搜索引擎"""

    def __init__(self):
        self.ids: List[int] = []
        self.name_vectors: np.ndarray = None
        self.spec_vectors: np.ndarray = None
        self.units: List[str] = []
        logger.info("NumpySearcher 初始化成功")

    def load_vectors(self, ids: List[int], name_vectors: np.ndarray,
                     spec_vectors: np.ndarray, units: List[str]):
        """加载向量数据"""
        self.ids = ids
        self.name_vectors = name_vectors
        self.spec_vectors = spec_vectors
        self.units = units
        logger.info(f"加载 {len(ids)} 个向量到搜索引擎")

    def normalize_vectors(self, vectors: np.ndarray) -> np.ndarray:
        """归一化向量"""
        if vectors is None or len(vectors) == 0:
            return vectors

        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        return vectors / np.where(norms == 0, 1, norms)

    def search_similar(self, query_names: List[str], query_specs: List[str],
                       query_units: List[str], top_k: int = 5) -> List[List[dict]]:
        """
        批量搜索相似项

        Returns:
            List[List[dict]]: 每个查询的top_k结果列表
        """
        if len(self.ids) == 0:
            logger.warning("搜索引擎中没有加载任何向量")
            return [[] for _ in range(len(query_names))]

        start_time = datetime.now()

        try:
            # 归一化查询向量
            query_name_norm = self.normalize_vectors(query_names)
            query_spec_norm = self.normalize_vectors(query_specs)

            # 归一化数据库向量
            db_name_norm = self.normalize_vectors(self.name_vectors)
            db_spec_norm = self.normalize_vectors(self.spec_vectors)

            # 计算相似度矩阵
            name_similarity = query_name_norm @ db_name_norm.T
            spec_similarity = query_spec_norm @ db_spec_norm.T

            results = []
            for i in range(len(query_names)):
                # 应用单位过滤
                unit_mask = np.array([1 if u == query_units[i] else 0
                                      for u in self.units], dtype=bool)

                # 应用阈值过滤
                name_mask = name_similarity[i] >= settings.NAME_SIMILARITY_THRESHOLD
                spec_mask = spec_similarity[i] >= settings.SPEC_SIMILARITY_THRESHOLD

                valid_mask = name_mask & spec_mask & unit_mask
                valid_indices = np.where(valid_mask)[0]

                if len(valid_indices) == 0:
                    results.append([])
                    continue

                # 计算综合分数
                combined_scores = (
                        settings.NAME_SIM_WEIGHT * name_similarity[i, valid_indices] +
                        settings.SPEC_SIM_WEIGHT * spec_similarity[i, valid_indices]
                )

                # 获取top_k结果
                top_indices = valid_indices[np.argsort(combined_scores)[-top_k:][::-1]]

                query_results = []
                for idx in top_indices:
                    query_results.append({
                        'id': self.ids[idx],
                        'similarity_score': float(combined_scores[np.where(valid_indices == idx)[0][0]]),
                        'name_similarity': float(name_similarity[i, idx]),
                        'spec_similarity': float(spec_similarity[i, idx])
                    })

                results.append(query_results)

            elapsed = (datetime.now() - start_time).total_seconds()
            logger.info(f"搜索完成，处理 {len(query_names)} 个查询，耗时 {elapsed:.3f}秒")

            return results

        except Exception as e:
            logger.error(f"搜索过程中发生错误: {e}")
            return [[] for _ in range(len(query_names))]

    def batch_performance_test(self, batch_sizes: List[int] = None):
        """批量性能测试"""
        if batch_sizes is None:
            batch_sizes = [1, 10, 100, 1000]

        if len(self.ids) == 0:
            logger.warning("无法进行性能测试：没有加载向量")
            return []

        results = []
        for batch_size in batch_sizes:
            # 生成测试数据
            vector_dim = self.name_vectors.shape[1]
            test_names = np.random.randn(batch_size, vector_dim)
            test_specs = np.random.randn(batch_size, vector_dim)
            test_units = [self.units[0]] * batch_size if self.units else ["个"] * batch_size

            start_time = datetime.now()
            self.search_similar(test_names, test_specs, test_units, top_k=5)
            elapsed = (datetime.now() - start_time).total_seconds()

            results.append({
                'batch_size': batch_size,
                'time_seconds': elapsed,
                'qps': batch_size / elapsed if elapsed > 0 else 0
            })

            logger.info(f"Batch size {batch_size}: {elapsed:.3f}s, QPS: {batch_size / elapsed:.1f}")

        return results


# 测试代码
if __name__ == "__main__":
    print("测试 NumpySearcher...")

    searcher = NumpySearcher()

    # 创建测试数据
    test_ids = [1, 2, 3, 4, 5]
    test_name_vecs = np.random.randn(5, 768)  # 5个样本，768维
    test_spec_vecs = np.random.randn(5, 768)
    test_units = ["个", "件", "台", "个", "件"]

    # 加载向量
    searcher.load_vectors(test_ids, test_name_vecs, test_spec_vecs, test_units)
    print("✓ 向量加载成功")

    # 创建查询数据
    query_names = np.random.randn(2, 768)  # 2个查询
    query_specs = np.random.randn(2, 768)
    query_units = ["个", "件"]

    # 执行搜索
    results = searcher.search_similar(query_names, query_specs, query_units, top_k=3)
    print("✓ 搜索完成")

    for i, result in enumerate(results):
        print(f"查询 {i + 1} 的结果:")
        for item in result:
            print(f"  ID: {item['id']}, 相似度: {item['similarity_score']:.3f}")

    # 性能测试
    perf_results = searcher.batch_performance_test([1, 2, 5])
    print("✓ 性能测试完成")

    for result in perf_results:
        print(f"Batch size {result['batch_size']}: {result['time_seconds']:.3f}s, QPS: {result['qps']:.1f}")

    print("所有测试完成！")